import time
from typing import List

import torch
from fastapi import FastAPI, Depends, HTTPException, File, Form, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from FlagEmbedding import BGEM3FlagModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# 导入配置
from app.core.config import settings
# 导入数据库客户端
from app.core.database import redis_client, chromadb_client
# 导入安全验证
from app.utils.auth_utils import verify_api_key
# 导入服务
from app.services.document_service import document_service
from app.services.search_service import init_search_service
from app.services.ai_service import ai_service
# 导入Pydantic模型
from app.model.schemas import SimpleQuery, Query, Document, DocumentBatch, Rerank, ModelRequest
from app.utils.decorators import timing_decorator, async_timing_decorator

# 初始化模型
print("正在加载BGE-M3模型...")
start_time = time.time()
model = BGEM3FlagModel(settings.BGE_M3_MODEL_PATH, use_fp16=True,
                       device="cuda:0" if torch.cuda.is_available() else "cpu")
end_time = time.time()
print(f"BGE-M3模型加载完成，耗时: {end_time - start_time:.2f}秒")

print("正在加载重排序模型...")
start_time = time.time()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
rerank_tokenizer = AutoTokenizer.from_pretrained(settings.BGE_RERANKER_MODEL_PATH)
rerank_model = AutoModelForSequenceClassification.from_pretrained(settings.BGE_RERANKER_MODEL_PATH)
rerank_model = rerank_model.half().to(device) if torch.cuda.is_available() else rerank_model.to(device)
end_time = time.time()
print(f"重排序模型加载完成，耗时: {end_time - start_time:.2f}秒")

# 初始化搜索服务
init_search_service(model, rerank_model, rerank_tokenizer, redis_client, chromadb_client)

from app.services.search_service import search_service

# FastAPI应用
app = FastAPI(title="RAG 系统 API", version="1.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["*"]
)


# 路由定义

@app.get("/")
async def root(api_key: str = Depends(verify_api_key)):
    return {"message": "RAG 系统 API 服务正常运行中"}


@app.post("/add_document", dependencies=[Depends(verify_api_key)])
async def add_document(doc: Document):
    try:
        result = await document_service.add_document(doc, model)
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/add_documents", dependencies=[Depends(verify_api_key)])
async def add_documents(batch: DocumentBatch):
    try:
        result = await document_service.add_documents(batch, model)
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.delete("/delete_document", dependencies=[Depends(verify_api_key)])
async def delete_document(doc_id: str, collection_name: str = "my_knowledge"):
    try:
        result = await document_service.delete_document(doc_id, collection_name)
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/get_document", dependencies=[Depends(verify_api_key)])
async def get_document(doc_id: str, collection_name: str = "my_knowledge"):
    try:
        result = await document_service.get_document(doc_id, collection_name)
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/upload", dependencies=[Depends(verify_api_key)])
@async_timing_decorator
async def upload(
        file: List[UploadFile] = File(...),
        chunk_size: int = Form(500),
        chunk_overlap: int = Form(100),
        collection_name: str = Form("my_knowledge")
):
    try:
        # 验证文件
        validation_result = await document_service.validate_files(file)
        if validation_result:
            return validation_result

        # 入库处理
        result = await document_service.embed_and_store_files(file, chunk_size, chunk_overlap, model, collection_name)
        return result
    except Exception as e:
        print(e)
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/collection_info", dependencies=[Depends(verify_api_key)])
async def collection_info(collection_name: str = None):
    try:
        result = await document_service.get_collection_info(collection_name)
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.delete("/delete_collection", dependencies=[Depends(verify_api_key)])
async def delete_collection(collection_name: str):
    try:
        result = await document_service.delete_collection(collection_name)
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/list_documents", dependencies=[Depends(verify_api_key)])
async def list_documents(
        collection_name: str = "my_knowledge",
        start: str = None,
        limit: int = 5
):
    try:
        result = await document_service.list_documents(collection_name, start, limit)
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/simple_query", dependencies=[Depends(verify_api_key)])
@async_timing_decorator
async def simple_query(query: SimpleQuery):
    try:
        results = search_service.hybrid_search(
            query.keyword,
            query.topK,
            query.collectionName,
            query.searchMethod,
            query.rerankWeight,
            query.scoreThreshold,
            query.rerankMethod,
        )
        return {"results": results}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/query", dependencies=[Depends(verify_api_key)])
@async_timing_decorator
async def query(query: Query):
    try:
        results = search_service.hybrid_search(
            query.keyword,
            query.topK,
            query.collectionName,
            query.searchMethod,
            query.rerankWeight,
            query.scoreThreshold,
            query.rerankMethod,
        )
        return {"results": results}
    except Exception as e:
        print(e)
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/rerank", dependencies=[Depends(verify_api_key)])
@async_timing_decorator
async def rerank(req: Rerank):
    try:
        return search_service.bge_rerank_batch(req.query, req.documents)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# 生成回答相关路由
@app.post("/model", dependencies=[Depends(verify_api_key)])
@async_timing_decorator
async def model_route(request: ModelRequest):
    try:
        return await ai_service.qwen3_32b(request)
    except Exception as e:
        print(e)
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/generate", dependencies=[Depends(verify_api_key)])
@async_timing_decorator
async def generate(query: Query):
    try:
        hybrid_search_results = search_service.hybrid_search(
            query.keyword,
            query.topK,
            query.collectionName,
            query.searchMethod,
            query.rerankWeight,
            query.scoreThreshold,
            query.rerankMethod
        )
        return await ai_service.generate_text(query.keyword, hybrid_search_results)
    except Exception as e:
        print(e)
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/stream_generate", dependencies=[Depends(verify_api_key)])
@async_timing_decorator
async def stream_generate(query: Query):
    try:
        hybrid_search_results = search_service.hybrid_search(
            query.keyword,
            query.topK,
            query.collectionName,
            query.searchMethod,
            query.rerankWeight,
            query.scoreThreshold,
            query.rerankMethod,
        )
        return await ai_service.stream_generate(query.keyword, hybrid_search_results, query.thinking)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# docker build -t rag-system .
# docker run -d -p 8000:8000 rag-system
# 前端部署在nginx上，可以通过http://localhost访问
# 数据库使用chromadb和redis，需要确保服务正常运行
# 模型需要预下载到指定路径，或通过环境变量指定路径
# 可以通过修改.env文件配置各种参数
# 更多使用说明，请查看README.md
